from __future__ import with_statement
import gc
import os
import sys
import thread
import itertools
import _hptrace
from types import FunctionType, BuiltinFunctionType
from contextlib import contextmanager


HPTracerError = _hptrace.error

class ThreadAlreadyTraced(HPTracerError):
    pass


#------------------------------------------------------------------------------ 
# marking function attributes
#------------------------------------------------------------------------------ 

def _get_code(func):
    if hasattr(func, "im_func"):
        func = func.im_func
    if hasattr(func, "func_code"):
        func = func.func_code
    return func

def _set_flag(func, flag):
    # python will not allow for unknown flags... pitty
    if isinstance(func, BuiltinFunctionType):
        pass
        #_hptrace._set_builtin_flags(func, flag)
    else:
        _hptrace._set_code_flags(_get_code(func), flag)
    return func

def _clear_flag(func, flag):
    if isinstance(func, BuiltinFunctionType):
        #_hptrace._clear_builtin_flags(func, flag)
        pass
    else:
        _hptrace._clear_code_flags(_get_code(func), flag)
    return func

def ignore_single(func):
    """signle level ignore. func can be code, builtin function, 
    function or method. can be applied as a decorator"""
    return _set_flag(func, _hptrace.CO_HPTRACE_IGNORED_SINGLE)

def ignore_recursive(func):
    """recursive ignore. func can be code, builtin function, 
    function or method. can be applied as a decorator"""
    return _set_flag(func, _hptrace.CO_HPTRACE_IGNORED_CHILDREN)

def ignore_whole(func):
    """single and recursive ignore. func can be code, builtin function, 
    function or method. can be applied as a decorator"""
    return _set_flag(func, _hptrace.CO_HPTRACE_IGNORED_WHOLE)

ignored = ignore_recursive

def unignore(func):
    return _clear_flag(func, _hptrace.CO_HPTRACE_IGNORED_WHOLE)

def detailed(func):
    """func can be code, function or method. can be applied as a decorator"""
    _hptrace._set_code_flags(_get_code(func), _hptrace.CO_HPTRACE_DETAILED)
    return func

def undetailed(func):
    """func can be code, function or method. can be applied as a decorator"""
    _hptrace._clear_code_flags(_get_code(func), _hptrace.CO_HPTRACE_DETAILED)
    return func

def ignore_all_by_predicate(codepred, bltpred, ignorefunc = ignore_recursive):
    for obj in gc.get_objects():
        if isinstance(obj, FunctionType) and codepred(obj):
            ignorefunc(obj)
        elif isinstance(obj, BuiltinFunctionType) and bltpred(obj):
            ignorefunc(obj)

#def ignore_path(path, deep = True, ignorefunc = ignore_recursive):
#    if deep:
#        def codepred(obj, path = path):
#            return obj.co_filename.startswith(path)
#        def bltpred(obj, modname = mod.__name__):
#            return obj.__module__ and obj.__module__.startswith(modname)
#    else:
#        def codepred(obj, fn = path):
#            return obj.co_filename.startswith(fn)
#        def bltpred(obj, modname = mod.__name__):
#            return obj.__module__ and obj.__module__.startswith(modname)
#
#    ignore_all_by_predicate(codepred, bltpred, ignorefunc = ignorefunc)
    
def ignore_module(mod, ignorefunc = ignore_recursive):
    fn = mod.__file__[:mod.__file__.rfind(".")]
    def codepred(obj, fn = fn):
        return obj.co_filename.startswith(fn)
    def bltpred(obj, modname = mod.__name__):
        return obj.__module__ == modname
    ignore_all_by_predicate(codepred, bltpred, ignorefunc = ignorefunc)

def ignore_package(mod, ignorefunc = ignore_recursive):
    if hasattr(mod, "__path__"):
        fn = mod.__path__[0]
    elif hasattr(mod, "__file__"):
        fn = mod.__file__.rsplit(".", 1)[0]
    else:
        fn = None
    if fn:
        def funcpred(obj, fn = fn):
            return obj.func_code.co_filename.startswith(fn)
    else:
        def funcpred(obj):
            return False
    def bltpred(obj, modname = mod.__name__):
        return obj.__module__ and obj.__module__.startswith(modname)
    ignore_all_by_predicate(funcpred, bltpred, ignorefunc = ignorefunc)


#------------------------------------------------------------------------------ 
# tracing and logging
#------------------------------------------------------------------------------ 

log = _hptrace.log

_thread_counter = itertools.count()
_per_thread = thread._local()
_thread_counter = itertools.count()

_orig_start_new_thread = thread.start_new_thread
_orig_start_new = thread.start_new

def _thread_wrapper(func, args, kwargs):
    with traced(g_fileprefix, threads = False, maxdepth = g_maxdepth):
        return func(*args, **kwargs)

def _start_new_thread(func, args, kwargs = {}):
    return _orig_start_new_thread(_thread_wrapper, (func, args, kwargs))

def load_config_file(filename = None):
    if filename is None:
        import __main__
        base = os.path.dirname(__main__.__file__)
        filename = os.path.join(base, "hptrace-config.py")
        if not os.path.exists(filename):
            return
    execfile(filename, {})

@contextmanager
def traced(fileprefix, threads = False, maxdepth = None, load_config = True):
    global g_fileprefix, g_maxdepth
    
    if getattr(_per_thread, "being_traced", False):
        raise ThreadAlreadyTraced("this thread already has an active tracer")
    
    restore_start_new = False
    if threads and thread.start_new_thread != _start_new_thread:
        thread.start_new = thread.start_new_thread = _start_new_thread
        restore_start_new = True
        g_fileprefix = fileprefix
        g_maxdepth = maxdepth
    
    if load_config:
        load_config_file()
    
    if maxdepth is None:
        maxdepth = sys.getrecursionlimit()
    fileprefix = fileprefix + "-%s" % (_thread_counter.next(),)
    tracer = _hptrace.HPTracer(fileprefix + ".cp", fileprefix + ".values", 
        fileprefix + ".trace", maxdepth)
    
    _per_thread.being_traced = True  
    try:
        tracer.start()
        yield
    finally:
        try:
            tracer.stop()
        finally:
            _per_thread.being_traced = False
            if restore_start_new:
                thread.start_new = _orig_start_new
                thread.start_new_thread = _orig_start_new_thread
                del g_fileprefix, g_maxdepth


if __name__ == "__main__":
    #------------------------------------------------------------------------- 
    # command line tool: trace python scripts
    #------------------------------------------------------------------------- 
    import time
    import sys
    from getopt import getopt, GetoptError

    def usage():
        print "usage: hptrace [options] script [arguments...]"
        print "where script is the python script filename to run, followed by the"
        print "arguments for that script"
        print 
        print "the options are:"
        print "    -h, --help       - print this help message and quit"
        print "    -f fileprefix    - the file prefix to use"
        print "    -d depth         - max stack depth (integer)"
        print "    -t               - trace threads too"
        print "    -c filename      - specify the config file to load, instead of"
        print "                       the default one (hptrace-config.py)"
        print "    --noconfig       - do not load default config file, even if it"
        print "                       exists"
        sys.exit()
    
    try:
        opts, args = getopt(sys.argv[1:], "f:td:h:c:", ["help", "noconfig"])
    except GetoptError, ex:
        print ex
        print
        usage()
    
    opts = dict(opts)
    if "-h" in opts or "--help" in opts or not args:
        usage()

    program = args[0]
    default_fileprefix = "%s-%s" % (program, time.strftime("%Y%m%d%H%M%S"))
    fileprefix = opts.get("-f", default_fileprefix)
    threads = "-t" in opts
    configfile = opts.get("-c", None)
    noconfig = "--noconfig" in opts
    if "-d" in opts:
        maxdepth = int(opts["-d"])
    else:
        maxdepth = None
    
    if not noconfig:
        load_ignore_file(configfile)
    
    with traced(fileprefix, threads = threads, maxdepth = maxdepth):
        sys.argv = args
        execfile(program, {"__name__" : "__main__"})





